from flask import Flask, request, Response, jsonify
from flask_cors import CORS
import requests
import json
import base64
import sys

app = Flask(__name__)
CORS(app)  # 允许跨域

VLLM_URL = "http://localhost:32770/v1/chat/completions"
MODEL_NAME = "/models/Qwen2.5-VL-7B-Instruct-AWQ"

@app.route('/chat', methods=['POST'])
def chat():
    data = request.get_json()
    user_msg = data.get('message', '')
    history = data.get('history', [])
    multimodal = data.get('multimodal', None)  # 多模态内容

    # 构建消息历史
    messages = []
    for msg in history:
        if msg.get('role') in ('user', 'ai'):
            role = 'user' if msg['role'] == 'user' else 'assistant'
            content = msg.get('content', '')
            if isinstance(content, list):
                # 处理多模态历史消息
                multimodal_content = []
                for item in content:
                    if isinstance(item, dict):
                        if item.get('type') == 'image_url' and 'image_url' in item:
                            # 确保图片URL格式正确
                            image_url = item['image_url'].get('url', '')
                            if image_url.startswith('data:'):
                                multimodal_content.append({
                                    "type": "image_url",
                                    "image_url": {"url": image_url}
                                })
                        elif item.get('type') == 'text':
                            multimodal_content.append(item)
                if multimodal_content:  # 只有当有内容时才添加消息
                    messages.append({"role": role, "content": multimodal_content})
            else:
                messages.append({"role": role, "content": [{"type": "text", "text": content}]})

    # 构建当前用户消息
    current_message = {"role": "user", "content": []}
    
    # 添加图片（如果有）
    if multimodal and multimodal.get("data"):
        image_data = multimodal['data']
        # 确保base64数据格式正确
        if not image_data.startswith('data:'):
            image_data = f"data:{multimodal.get('mimetype', 'image/jpeg')};base64,{image_data}"
        current_message["content"].append({
            "type": "image_url",
            "image_url": {
                "url": image_data
            }
        })
    
    # 添加文本（如果有）
    if user_msg:
        current_message["content"].append({
            "type": "text",
            "text": user_msg
        })
    
    # 如果current_message包含image_url, 且messages里面包含image_url消息，去掉message里面对应消息对象
    if current_message["content"] and any(item.get("type") == "image_url" for item in current_message["content"]):
        # 获取当前消息中的图片URL
        current_image_urls = [
            item["image_url"]["url"] 
            for item in current_message["content"] 
            if item.get("type") == "image_url" and "image_url" in item
        ]
        
        # 遍历历史消息，移除所有包含图片的消息
        filtered_messages = []
        for msg in messages:
            if msg["role"] == "user" and isinstance(msg.get("content"), list):
                # 检查消息中是否包含图片
                has_image = any(
                    item.get("type") == "image_url" and "image_url" in item 
                    for item in msg["content"]
                )
                # 如果消息不包含图片，保留该消息
                if not has_image:
                    filtered_messages.append(msg)
            else:
                filtered_messages.append(msg)
        
        messages = filtered_messages

    messages.append(current_message)

    def generate():
        headers = {"Content-Type": "application/json"}
        payload = {
            "model": MODEL_NAME,
            "messages": messages,
            "temperature": 0.7,
            "max_tokens": 512,
            "top_p": 0.8,
            "stream": True
        }
        print("[VLLM Request Payload]", file=sys.stderr)
        print(json.dumps(payload, ensure_ascii=False, indent=2), file=sys.stderr)
        try:
            with requests.post(VLLM_URL, headers=headers, data=json.dumps(payload), stream=True) as resp:
                print(f"[VLLM Response Status] {resp.status_code}", file=sys.stderr)
                for line in resp.iter_lines():
                    if line:
                        try:
                            print(f"[VLLM Response Line] {line}", file=sys.stderr)
                            # 处理SSE格式的数据
                            line = line.decode('utf-8')
                            if line.startswith('data: '):
                                line = line[6:]  # 移除 'data: ' 前缀
                            if line.strip() == '[DONE]':
                                continue
                            data = json.loads(line)
                            if 'choices' in data and len(data['choices']) > 0:
                                delta = data['choices'][0].get('delta', {})
                                if 'content' in delta:
                                    yield delta['content']
                        except Exception as e:
                            print(f"[VLLM Parse Error] {e}", file=sys.stderr)
                            continue
        except Exception as e:
            print(f"[VLLM Request Error] {e}", file=sys.stderr)
            yield "[后端请求VLLM出错，请检查模型和服务是否正常]"

    return Response(generate(), mimetype='text/plain')

@app.route('/upload', methods=['POST'])
def upload():
    file = request.files.get('file')
    if not file:
        return jsonify({'error': 'No file uploaded'}), 400
    mimetype = file.mimetype
    filetype = 'image' if mimetype.startswith('image/') else 'video' if mimetype.startswith('video/') else 'file'
    b64 = base64.b64encode(file.read()).decode('utf-8')
    return jsonify({
        'type': filetype,
        'data': b64,
        'mimetype': mimetype
    })

if __name__ == '__main__':
    app.run(host='127.0.0.1', port=5000, threaded=True)
